import numpy as np
import torch
from scipy.spatial.transform import Rotation as R
from numprint import p

def debug_quaternion_issue():
    real_pose = np.array([0.064, 0.007, 0.198, 2.084, 1.340, 2.183, 0.000])
    calc_pose = np.array([[0.055, 0.001, 0.203, 0.631, 0.406, 0.661, 0.0]])
    
    print("=== 问题分析 ===")
    p.print("real_pose", real_pose)
    p.print("calc_pose", calc_pose)
    
    real_quat = real_pose[3:7]
    calc_quat = calc_pose[0, 3:7]
    
    print("\n=== 四元数分析 ===")
    p.print("real_quat", real_quat)
    p.print("calc_quat", calc_quat)
    print(f"real_quat norm: {np.linalg.norm(real_quat)}")
    print(f"calc_quat norm: {np.linalg.norm(calc_quat)}")
    
    print("\n=== 检查是否是欧拉角 ===")
    print("real_quat 看起来像欧拉角 (弧度):", real_quat)
    print("转换为度:", np.degrees(real_quat))
    
    print("\n=== 尝试将欧拉角转换为四元数 ===")
    try:
        real_rot_from_euler = R.from_euler('xyz', real_quat)
        real_quat_from_euler = real_rot_from_euler.as_quat()
        p.print("real_quat_from_euler", real_quat_from_euler)
        print(f"real_quat_from_euler norm: {np.linalg.norm(real_quat_from_euler)}")
    except Exception as e:
        print(f"欧拉角转换失败: {e}")
    
    print("\n=== 检查原始代码中的四元数处理 ===")
    print("在 franka_fk.py 中，fk_qxyzw 方法:")
    print("quaternion_reordered = torch.cat([kin_state.ee_quaternion[:, 1:], kin_state.ee_quaternion[:, 0:1]], dim=-1)")
    print("这意味着从 [w,x,y,z] 重新排序为 [x,y,z,w]")
    
    print("\n=== 检查 calc_quat 是否也需要重新排序 ===")
    calc_quat_wxyz = np.array([calc_quat[3], calc_quat[0], calc_quat[1], calc_quat[2]])
    p.print("calc_quat_wxyz", calc_quat_wxyz)
    print(f"calc_quat_wxyz norm: {np.linalg.norm(calc_quat_wxyz)}")
    
    print("\n=== 正确的变换计算 ===")
    if np.linalg.norm(real_quat) > 2:
        print("real_quat 看起来是欧拉角，先转换为四元数")
        real_rot = R.from_euler('xyz', real_quat)
        real_quat_correct = real_rot.as_quat()
    else:
        real_quat_correct = real_quat
    
    calc_rot = R.from_quat(calc_quat)
    real_rot_correct = R.from_quat(real_quat_correct)
    
    p.print("real_quat_correct", real_quat_correct)
    p.print("calc_quat", calc_quat)
    
    quat_offset = (calc_rot.inv() * real_rot_correct).as_quat()
    p.print("quat_offset", quat_offset)
    
    transformed_quat = (calc_rot * R.from_quat(quat_offset)).as_quat()
    p.print("transformed_quat", transformed_quat)
    p.print("real_quat_correct", real_quat_correct)
    
    print(f"变换是否正确: {np.allclose(transformed_quat, real_quat_correct, atol=1e-6)}")

if __name__ == "__main__":
    debug_quaternion_issue()
